import torch.nn as nn
from ..builder import BACKBONES


@BACKBONES.register_module()
class MyNet(nn.Module):
    def __init__(self, in_c, out_c, pool: int = 1):
        super(MyNet, self).__init__()
        self.cov1 = nn.Conv2d(in_channels=in_c, out_channels=in_c, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.bn1 = nn.BatchNorm2d(in_c, eps=1e-3, momentum=0.01)
        self.cov2 = nn.Conv2d(in_channels=in_c, out_channels=out_c, kernel_size=(3, 3), stride=(pool, pool),
                              padding=(1, 1))
        self.bn2 = nn.BatchNorm2d(in_c, eps=1e-3, momentum=0.01)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x1 = self.cov1(x)
        x1 = self.bn1(x1)
        x1 = self.relu(x1)
        x1 = self.cov2(x1)
        x1 = self.bn2(x1)
        x1 = self.relu(x1)
        return x1
